import logging
import os

import cv2
import numpy as np
from tflite_runtime import interpreter as tflite

# face tf-lite model
FACE_MODEL = "./model/face/facenet.tflite"

# valid data size
VALID_DATA_SIZE = 3


class ImageClass:

    def __init__(self, name, image_paths):
        self.name = name
        self.image_paths = image_paths

    def __str__(self):
        return self.name + ', ' + str(len(self.image_paths)) + ' images'

    def __len__(self):
        return len(self.image_paths)


def train_dataset(root_path, alter_paths=None):
    if alter_paths is None:
        dataset = get_all_data(root_path)
    else:
        dataset = get_selected_data(root_path, alter_paths)
    if len(dataset) == 0:
        return None, None
    features = []
    names = []
    interpreter = tflite.Interpreter(model_path=FACE_MODEL)
    input_details = interpreter.get_input_details()
    logging.info(input_details)
    output_details = interpreter.get_output_details()
    interpreter.allocate_tensors()
    for i in range(len(dataset)):
        paths = dataset[i].image_paths
        name = dataset[i].name
        for path in paths:
            face = cv2.imread(path)
            face_id = cv2.resize(face, (160, 160))
            image_array = np.asarray(face_id)
            normalized_image_array = (image_array.astype(np.float32) / 127.0) - 1
            interpreter.set_tensor(input_details[0]['index'], [normalized_image_array])
            interpreter.invoke()
            feature = interpreter.get_tensor(output_details[0]['index'])
            features.append(feature)
            names.append(name)
    return features, names


def get_all_data(root_path):
    path_exp = os.path.expanduser(root_path)
    all_paths = sorted([path for path in os.listdir(path_exp)
                        if os.path.isdir(os.path.join(path_exp, path))])
    return get_selected_data(root_path, all_paths)


def get_selected_data(root_path, paths: list):
    dataset = []
    path_exp = os.path.expanduser(root_path)
    for i in range(len(paths)):
        class_name = paths[i]
        face_dir = os.path.join(path_exp, class_name)
        image_paths = get_image_paths(face_dir)
        if len(image_paths) > 0:
            dataset.append(ImageClass(class_name, image_paths))
    return dataset


def get_image_paths(face_dir):
    image_paths = []
    if os.path.isdir(face_dir):
        images = os.listdir(face_dir)
        image_paths = [os.path.join(face_dir, img) for img in images]
    return image_paths[:VALID_DATA_SIZE]
